#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
TypeInferencePass - 类型推导阶段

职责：
1. 只推导不明确的类型（没有用户注解的节点）
2. 不处理已有注解类型的节点
3. 不进行类型匹配检查（留给TypeCheckingPass）
4. 将推导结果直接存储到IR节点的extension_params中
5. 如果启用严格类型检查，验证所有变量、函数参数和返回类型都有类型注解
"""

from typing import Any, Dict, Optional, Set, List

from mcpy.ir.module import IRModule
from .base import BasePass, PassStage, PassContext
from ..ir.base import IRNode
from ..type_inference import infer_type, InferenceContext, InferenceMode
from ..ir.declarations import IRImportNode, IRFunctionNode, IRMethodNode
from ..ir.expressions import IRName, IRConstant
from ..ir.statements import IRAssign, IRExprStmt, IRTryExcept
from ..symbol_meta import TypeInfo


class StrictTypingError(Exception):
    """严格类型检查错误"""

    def __init__(self, errors: List[str]):
        self.errors = errors
        super().__init__(f"严格类型检查失败，共 {len(errors)} 个错误")


SUPPORTED_NODE_TYPES = {
    "assign",  # 赋值语句
    "expr_stmt",  # 表达式语句
    "call_expr",  # 函数调用表达式
    "name_expr",  # 名称表达式
    "constant",  # 常量表达式
    "bin_op",  # 二元操作表达式
    "unary_op",  # 一元操作表达式
    "compare",  # 比较操作表达式
    "attribute",  # 属性访问表达式
    "subscript",  # 下标访问表达式
    "property",  # 属性节点
    "function",  # 函数节点
    "class",  # 类节点
}

# 临时变量名（在严格类型检查中会被忽略）
TEMP_VAR_NAMES = {"_", "i", "j", "temp", "tmp", "result"}


class TypeInferencePass(BasePass):
    """类型推导Pass

    专门负责推导不明确的类型，不干预已有注解的类型。
    """

    name = "type_inference"
    stage = PassStage.TYPE_ANALYSIS

    def __init__(self, config: Optional[Dict[str, Any]] = None):
        super().__init__()
        self.config = config or {}
        self.processed_nodes: Set[int] = set()  # 使用对象ID而不是对象本身
        self.inference_context: Optional[InferenceContext] = None

        # 严格类型检查配置
        self.strict_typing = self.config.get("strict_typing", False)
        self.type_errors: List[str] = []  # 收集类型注解错误

    def run(self, ir_module: IRModule, context: PassContext) -> IRNode:
        """运行类型推导

        Args:
            ir_module: IR模块节点

        Returns:
            处理后的IR模块节点
        """

        # 清空已处理节点集合和错误列表
        self.processed_nodes.clear()
        self.type_errors.clear()

        self.inference_context = InferenceContext(
            mode=InferenceMode.LOCAL, compiler=context.compiler
        )

        # 先 resolve 所有 import 语句，确保所有符号都已导入
        self._resolve_all_imports(ir_module, context.compiler)

        # 遍历所有节点进行类型推导
        self._infer_types(ir_module)

        # 如果启用严格类型检查，验证错误并抛出
        if self.strict_typing and self.type_errors:
            raise StrictTypingError(self.type_errors)

        return ir_module

    def _resolve_all_imports(self, ir_module: IRModule, compiler) -> None:
        """预先 resolve 所有 import 语句，确保符号在类型推导前可用

        Args:
            ir_module: IR 模块节点
            compiler: 编译器实例
        """

        for node in ir_module.get_child_nodes():
            if isinstance(node, IRImportNode):
                node.resolve_import(compiler)

    def _infer_types(self, node: IRNode) -> None:
        """推导模块中所有节点的类型

        Args:
            node: 当前处理的节点
        """
        # 避免重复处理
        node_id = id(node)
        if node_id in self.processed_nodes:
            return
        self.processed_nodes.add(node_id)

        # 如果是函数节点，设置 current_function

        if isinstance(node, (IRFunctionNode, IRMethodNode)):
            # 创建上下文，设置 current_function
            old_context = self.inference_context
            self.inference_context = InferenceContext(
                mode=old_context.mode if old_context else InferenceMode.LOCAL,
                error_messages=old_context.error_messages if old_context else [],
                current_class_type=(
                    old_context.current_class_type if old_context else None
                ),
                infer_stack=old_context.infer_stack.copy() if old_context else [],
                compiler=old_context.compiler if old_context else None,
                expected_type=None,
                current_function=node,
            )

        if self._should_infer_type(node):
            self._infer_node_type(node)

        # 严格类型检查
        if self.strict_typing:
            self._check_type_annotations(node)

        for child in node.get_child_nodes():
            self._infer_types(child)

        # 恢复上下文
        if isinstance(node, (IRFunctionNode, IRMethodNode)):
            self.inference_context = old_context

    def _should_infer_type(self, node: IRNode) -> bool:
        """判断是否需要推导类型

        Args:
            node: 要检查的节点

        Returns:
            如果需要推导类型返回True，否则返回False
        """
        # 赋值语句：仅在有类型注解时重新处理（支持双向推导）
        if node.node_type == "assign":
            # 检查是否有类型注解
            has_annotation = getattr(node, "annotation")
            if has_annotation:
                # 有注解：需要处理（支持双向推导和 Any 类型）
                return True
            # 无注解：检查是否已有类型
            if node.has_extension_param("inferred_type"):
                return False  # 已有类型，跳过
            return True  # 无类型，需要推导

        # 如果已经有推导类型，跳过
        if node.has_extension_param("inferred_type"):
            return False

        # 只处理需要类型推导的节点类型
        # return node.node_type in SUPPORTED_NODE_TYPES
        return True

    def _infer_node_type(self, node) -> None:
        """推导单个节点的类型 - 使用统一入口点，同时更新符号信息

        Args:
            node: 要推导类型的节点
        """
        if not self.inference_context:
            return

        inferred_type = infer_type(node, self.inference_context)
        node.set_extension_param("inferred_type", inferred_type, replace_existing=False)

        # 如果推导成功且节点有名称，更新符号的类型信息
        symbol_name = node.name
        if symbol_name:
            symbol = node.lookup_symbol(symbol_name)
            if symbol:
                symbol.set_type_info(inferred_type)

    # ===== 严格类型检查方法 =====

    def _check_type_annotations(self, node: IRNode) -> None:
        """检查节点的类型注解是否完整

        Args:
            node: 要检查的节点
        """
        # 检查模块级别变量
        if node.node_type == "assign" and self._is_module_level_assignment(node):
            self._check_variable_annotation(node)

        # 检查函数和方法
        elif node.node_type in ("function", "method"):
            self._check_function_annotations(node)

    def _is_module_level_assignment(self, node: IRNode) -> bool:
        """判断赋值是否在模块级别（不在函数/类内部）

        Args:
            node: 赋值节点

        Returns:
            True 如果是模块级别赋值
        """
        parent = node.parent
        while parent:
            # 如果父节点是函数或方法，说明是函数内部变量
            if parent.node_type in ("function", "method"):
                return False
            # 注意：类内部的类属性仍然算作模块级别（需要类型注解）
            parent = parent.parent
        return True

    def _check_variable_annotation(self, node: IRNode) -> None:
        """检查变量是否有类型注解

        Args:
            node: 赋值节点
        """

        # 确保是赋值节点
        if not isinstance(node, IRAssign):
            return

        # 如果已经有类型注解，跳过
        if node.annotation:
            return

        # 获取所有目标表达式
        target_exprs = node.get_target_exprs()
        for target in target_exprs:
            # 只检查简单变量名
            if isinstance(target, IRName):
                var_name = target.id

                # 忽略临时变量
                if self._is_temp_variable(var_name):
                    continue

                # 报告错误
                line_info = (
                    f"（位于第 {node.source_line} 行）" if node.source_line else ""
                )
                self.type_errors.append(f"变量 '{var_name}' 缺少类型注解{line_info}")

    def _check_function_annotations(self, node: IRNode) -> None:
        """检查函数参数和返回类型注解

        Args:
            node: 函数或方法节点
        """

        if not isinstance(node, (IRFunctionNode, IRMethodNode)):
            return

        func_name = node.name
        line_info = f"（位于第 {node.source_line} 行）" if node.source_line else ""

        # 检查参数类型注解
        for i, param in enumerate(node.params):
            # 跳过 self 参数（类方法的第一个参数）
            if isinstance(node, IRMethodNode) and i == 0 and param.name == "self":
                continue

            # 检查是否有类型注解
            if not param.type_annotation:
                self.type_errors.append(
                    f"函数 '{func_name}' 的参数 '{param.name}' 缺少类型注解{line_info}"
                )

        # 检查返回类型注解
        # __init__ 方法和 main 函数可以没有返回类型注解
        # 只包含 pass 语句的函数也可以没有返回类型注解
        if func_name not in ("__init__", "main") and not node.return_type:
            # 检查函数体是否只包含 pass 语句
            if not self._is_empty_function(node):
                self.type_errors.append(
                    f"函数 '{func_name}' 缺少返回类型注解{line_info}"
                )

    def _is_temp_variable(self, name: str) -> bool:
        """判断是否是临时变量

        Args:
            name: 变量名

        Returns:
            True 如果是临时变量
        """
        return name in TEMP_VAR_NAMES

    def _is_empty_function(self, node: IRNode) -> bool:
        """判断函数是否为空（只包含 pass 或 docstring）

        Args:
            node: 函数节点

        Returns:
            True 如果函数为空
        """

        if not isinstance(node, (IRFunctionNode, IRMethodNode)):
            return False

        # 获取函数体中的语句
        if not hasattr(node, "body") or not hasattr(node.body, "statements"):
            return True

        statements = node.body.statements

        # 如果没有语句，认为是空函数
        if not statements:
            return True

        # 如果只有一个语句，且该语句是表达式语句（可能是 pass 转换的）
        # 或者语句的表达式是 None 常量，则认为是空函数
        if len(statements) == 1:
            stmt = statements[0]
            if isinstance(stmt, IRExprStmt):

                if hasattr(stmt, "expr") and isinstance(stmt.expr, IRConstant):
                    if stmt.expr.value is None:
                        return True

        return False
